# -*- coding:utf-8 -*-
"""
作者：ChenTao
日期：2023年03月02日
"""

import numpy as np
import cv2
import math


class Pose_estimation:
    def __init__(self,
                 camera_matrix=np.array([[4161.221, 0, 1445.577], [0, 4161.221, 984.686], [0, 0, 1]]) / 4,
                 dist_coeffs=np.array([[-0.1065, 0.0793, -0.0002, -8.9263e-06, -0.0161]])):
        """
        姿态估计
        Args:
            camera_matrix: 相机内参
            dist_coeffs: 畸变矩阵

        """
        self.camera_matrix = camera_matrix
        self.dist_coeffs = dist_coeffs

    def P4p(self, pixel_points, world_points):
        """

        Args:
            pixel_points: 像素坐标
            world_points: 世界坐标

        Returns:两个转角

        """
        success, rotation_vector, translation_vector = cv2.solvePnP(world_points,
                                                                    pixel_points,
                                                                    self.camera_matrix,
                                                                    self.dist_coeffs,
                                                                    flags=cv2.SOLVEPNP_ITERATIVE)

        rotM = cv2.Rodrigues(rotation_vector)[0]
        position = -np.matrix(rotM).T * np.matrix(translation_vector)
        # print("file Vector:\n {0}".format(-np.matrix(rotM).T * np.matrix(translation_vector)))

        yaw = Radians2Angles(math.atan(position[0, 0] / position[2, 0]))
        pitch = Radians2Angles(math.atan(position[1, 0] / position[2, 0]))

        print(yaw, pitch)
        return pitch, yaw

    def dedistortion(self, src, mode):
        """
        去畸变
        Args:
            src: 图片
            mode: 模式，图像去畸变和像素点去畸变

        Returns: 去畸变后的像素坐标

        """
        if mode == 'image':
            dedistortion_image = cv2.undistort(src,
                                               self.camera_matrix,
                                               self.dist_coeffs,
                                               P=self.camera_matrix)
        elif mode == 'pixel':
            dedistortion_points = cv2.undistortPoints(src,
                                                      self.camera_matrix,
                                                      self.dist_coeffs,
                                                      P=self.camera_matrix)

        return dedistortion_points[0][0][0], dedistortion_points[0][0][1]

    def pixelToangle(self, x, y):
        """
        根据 2D像素点推出旋转角度
        Args:
            x: 坐标
            y: 坐标

        Returns: 转角

        """
        pitch_list, yaw_list = [], []  # 两个列表相同下标对应同一个目标

        if isinstance(x, list):
            x_list, y_list = x, y
            for x, y in zip(x_list, y_list):
                # x, y = self.dedistortion((x, y), mode='pixel')
                pitch = math.atan((x - self.camera_matrix[0][2]) / self.camera_matrix[0][0])
                yaw = math.atan((y - self.camera_matrix[1][2]) / self.camera_matrix[1][1])

                pitch_list.append(round(Radians2Angles(pitch), 1))
                yaw_list.append(round(Radians2Angles(yaw), 1))
        else:
            pitch = math.atan((x - self.camera_matrix[0][2]) / self.camera_matrix[0][0])
            yaw = math.atan((y - self.camera_matrix[1][2]) / self.camera_matrix[1][1])

            pitch_list.append(round(Radians2Angles(pitch), 1))
            yaw_list.append(round(Radians2Angles(yaw), 1))

        # print('排序前：{} {}'.format(pitch_list, yaw_list))
        pitch_list, yaw_list = zip(*sorted(zip(pitch_list, yaw_list)))  # 对两个关联列表排序
        # print('排序后：{} {}'.format(pitch_list, yaw_list))

        return pitch_list, yaw_list


def Radians2Angles(radians):
    """
    弧度制转化为角度制
    Args:
        radians: 弧度

    Returns: 角度

    """
    return radians / math.pi * 180
